import numpy as np

from src.ga.impl.tf import get_tf_output
from src.ga.impl.pytorch import get_torch_output
from src.ga.impl.caffe import get_caffe_output

import tensorflow as tf


def get_output(tensor, operator):
    if operator == 'avg_pool':
        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy()
    elif operator == 'avg_pool3d':
        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy()
    elif operator == 'bias_add':
        variable = tf.normal.random([np.asarray(tensor).shape[-1]], stddev=0.01).numpy()
        return get_tf_output(tensor, operator, variable).numpy(), get_torch_output(tensor, operator, variable).numpy()